from typing import Any, Dict, List
import chromadb
from elasticsearch import AsyncElasticsearch
from sqlalchemy.ext.asyncio import create_async_engine
from neo4j import AsyncGraphDatabase
from rank_bm25 import BM25Okapi
from ..base import BaseSearchStrategy, Document, SearchResult


class VectorSearchStrategy(BaseSearchStrategy):
    """向量搜索策略"""
    def __init__(self):
        self.client = None
        self.collection = None

    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化Chroma客户端"""
        self.client = chromadb.Client(config["path"])
        self.collection = self.client.create_collection(
            name="documents",
            embedding_function=config["embedding_model"]
        )

    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        required = ["path", "embedding_model"]
        return all(key in config for key in required)

    async def search(self, query: str, **kwargs) -> SearchResult:
        """执行向量搜索"""
        k = kwargs.get("k", 5)
        results = self.collection.query(
            query_texts=[query],
            n_results=k
        )
        
        documents = [
            Document(
                content=doc["text"],
                metadata=doc["metadata"],
                score=score
            )
            for doc, score in zip(results["documents"][0], results["distances"][0])
        ]
        
        return SearchResult(
            documents=documents,
            source="vector",
            metadata={"strategy": "chroma"}
        )


class TextSearchStrategy(BaseSearchStrategy):
    """文本搜索策略"""
    def __init__(self):
        self.client = None

    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化Elasticsearch客户端"""
        self.client = AsyncElasticsearch(
            hosts=[f"{config['host']}:{config['port']}"]
        )
        self.index = config["index"]

    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        required = ["host", "port", "index"]
        return all(key in config for key in required)

    async def search(self, query: str, **kwargs) -> SearchResult:
        """执行文本搜索"""
        k = kwargs.get("k", 5)
        response = await self.client.search(
            index=self.index,
            body={
                "query": {
                    "multi_match": {
                        "query": query,
                        "fields": ["content", "title"]
                    }
                },
                "size": k
            }
        )
        
        documents = [
            Document(
                content=hit["_source"]["content"],
                metadata=hit["_source"].get("metadata", {}),
                score=hit["_score"]
            )
            for hit in response["hits"]["hits"]
        ]
        
        return SearchResult(
            documents=documents,
            source="text",
            metadata={"strategy": "elasticsearch"}
        )


class SQLSearchStrategy(BaseSearchStrategy):
    """SQL搜索策略"""
    def __init__(self):
        self.engine = None

    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化SQL引擎"""
        self.engine = create_async_engine(config["url"])

    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        required = ["url"]
        return all(key in config for key in required)

    async def search(self, query: str, **kwargs) -> SearchResult:
        """执行SQL搜索"""
        # 这里需要实现NL2SQL转换
        # 示例实现
        sql = f"SELECT * FROM documents WHERE content LIKE '%{query}%' LIMIT 5"
        
        async with self.engine.connect() as conn:
            result = await conn.execute(sql)
            rows = await result.fetchall()
            
        documents = [
            Document(
                content=str(row["content"]),
                metadata={"table": row["table_name"]},
                score=1.0
            )
            for row in rows
        ]
        
        return SearchResult(
            documents=documents,
            source="sql",
            metadata={"strategy": "postgresql"}
        )


class GraphSearchStrategy(BaseSearchStrategy):
    """图搜索策略"""
    def __init__(self):
        self.driver = None

    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化Neo4j驱动"""
        self.driver = AsyncGraphDatabase.driver(
            config["url"],
            auth=(config["user"], config["password"])
        )

    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        required = ["url", "user", "password"]
        return all(key in config for key in required)

    async def search(self, query: str, **kwargs) -> SearchResult:
        """执行图搜索"""
        # 这里需要实现NL2Cypher转换
        # 示例实现
        cypher = f"""
        MATCH (n)
        WHERE n.content CONTAINS $query
        RETURN n
        LIMIT 5
        """
        
        async with self.driver.session() as session:
            result = await session.run(cypher, query=query)
            records = await result.fetch()
            
        documents = [
            Document(
                content=record["n"]["content"],
                metadata={"labels": list(record["n"].labels)},
                score=1.0
            )
            for record in records
        ]
        
        return SearchResult(
            documents=documents,
            source="graph",
            metadata={"strategy": "neo4j"}
        )


class BM25SearchStrategy(BaseSearchStrategy):
    """BM25搜索策略"""
    def __init__(self):
        self.bm25 = None
        self.corpus = None
        self.doc_mapping = None

    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化BM25模型"""
        self.corpus = config["corpus"]
        self.doc_mapping = config["doc_mapping"]
        tokenized_corpus = [doc.split() for doc in self.corpus]
        self.bm25 = BM25Okapi(tokenized_corpus)

    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        required = ["corpus", "doc_mapping"]
        return all(key in config for key in required)

    async def search(self, query: str, **kwargs) -> SearchResult:
        """执行BM25搜索"""
        k = kwargs.get("k", 5)
        tokenized_query = query.split()
        scores = self.bm25.get_scores(tokenized_query)
        
        # 获取前k个最相关文档
        top_k_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k]
        
        documents = [
            Document(
                content=self.corpus[idx],
                metadata=self.doc_mapping[idx],
                score=scores[idx]
            )
            for idx in top_k_indices
        ]
        
        return SearchResult(
            documents=documents,
            source="bm25",
            metadata={"strategy": "bm25"}
        ) 